import random
import hashlib

# 快速幂取模 a^b % c
def fastMod(a, b, c):
    result = 1
    # 重点是将指数b做二进制分解，遍历b的每一个二进制位
    while b!= 0:
        # 当前位置是非0，进行乘模运算
        if (b&1) == 1 :
            result = (result * a) % c
        b = b >> 1
        # 每轮不断更新为a^2 mod c, a^4 mod c, ...
        a = (a * a) % c
    return result

# Miller-Rabin 素性检测算法 test(n)
def test(n):

    # 小质数加速

    primeIn100 = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97]

    if n in primeIn100:
        return True
    
    for prime in primeIn100:
        if n % prime == 0:
            return False

    # 找到 k 和 q，满足n = 2^k * q
    k = 0
    q = n-1
    while q%2 == 0 :
        q = q // 2
        k = k + 1

    isPrime = True # 初始化假定 n 是质数

    # 进行10轮检验
    for i in range(10):
        # 随机选择一个a，满足1<a<n-1
        a = random.randint(2, n-2)

        # 条件1判断
        # 快速幂取模算法计算 a^q mod n是否为 1 或 n-1
        if(fastMod(a, q, n)==1 or fastMod(a, q, n)==n-1):
            continue # 是质数

        # 条件1不满足，进入条件2判断
        # 快速幂取模算法计算 a^((2^j)^q) mod n 是否为 n-1
        satisfy2 = False # 初始化 条件2 不满足
        for j in range(1, k):
            if(fastMod(a, pow(2, j)*q, n)==n-1):
                satisfy2 = True # 其中一次满足即可
        
        # 满足条件2
        if satisfy2 :
            continue

        # 条件1 和 条件2均不满足
        isPrime = False
    
    return isPrime

# 生成一个位数为bits大质数 p
def getLargePrimes(bits):
    p = random.getrandbits(bits) | (1 << (bits-1))
    # 直到找到一个素数p
    while not test(p):
        p = random.getrandbits(bits) | (1 << (bits-1))
    return p

# 欧几里得算法求两个数最大公因数
def gcd(a, b):
    # 保证 a>=b 
    if a < b: 
        a, b = b, a
    # 辗转相除法
    while a%b != 0:
        a, b = b, a%b
    return b


# 拓展的欧几里得算法 a*x + b*y = gcd(a, b)，求x，y，gcd(a, b)
def ext_gcd(a, b):
    change = False
    # 确保 a >= b
    if a < b:
        a, b = b, a
        change = True
    x1, y1, r1 = 1, 0, a
    x2, y2, r2 = 0, 1, b
    x, y, r = (x1 - (r1//r2)*x2), (y1 - (r1//r2)*y2), r1%r2
    while r != 0 :
        x1, y1, r1 = x2, y2, r2
        x2, y2, r2 = x, y, r
        x, y, r = (x1 - (r1//r2)*x2), (y1 - (r1//r2)*y2), r1%r2
    # 恢复
    if change :
        x2, y2, r2 = y2, x2, r2
    return x2, y2, r2


# 获得逆元 a * a^-1 = 1 mod b
def getInverse(a, b):
    x, y, r = ext_gcd(a, b)
    if r!=1 :
        print("获取逆元失败")
    else:
        return x % b


# 随机获取一个大素数的原根
def getGenerator(prime):
    # prime为质数, 如果g的阶为(p-1), 则g是一个生成元
    # 随机产生p满足1<g<prime
    g = random.randint(2, prime-1)
    # 验证g^(prime-1) mod prime是否为1
    while fastMod(g, prime-1, prime)!=1 :
        # 重新生成
        g = random.randint(2, prime-1)
    return g


# 生成公钥PK和私钥SK
def getPKSK():
    # 随机产生一个大素数
    p = getLargePrimes(512)

    # 获取p的原根
    g = getGenerator(p)

    # 随机生成一个私钥 x
    x = random.randint(2, p-1)

    # 根据私钥计算公钥 y
    y = fastMod(g, x, p)

    # 全部私钥
    SK = [p, g, x]
    PK = [p, g, y]

    return PK, SK


# 签名算法
def sign(m, SK):
    p = SK[0] # 大质数
    g = SK[1] # 生成元
    x = SK[2] # 私钥
    
    # 选择一个随机数k 满足1<=k<=p-1 且 k和p-1互质
    k = random.randint(1, p-1)
    while gcd(k, p-1)!=1 :
        k = random.randint(1, p-1)

    # 得到k模q-1的逆元 k_inverse
    k_inverse = getInverse(k, p-1)
    
    # 计算H(m)
    H_m = eval("0x" + hashlib.sha256(m.encode('utf-8')).hexdigest())

    # 计算r
    r = fastMod(g, k, p)

    # 计算s
    s = (k_inverse * (H_m - x * r)) % (p-1)

    print("本轮选取的随机数: ")
    print("k: " + str(k))
    print()
    
    # 返回(r, s)
    return r, s


# 验证算法
def verify(m, r, s, PK):
    p = PK[0]
    g = PK[1]
    y = PK[2]

    # 计算H(m)
    H_m = eval("0x" + hashlib.sha256(m.encode('utf-8')).hexdigest())

    # 计算第一个模数
    mod1 = (fastMod(y, r, p) * fastMod(r, s, p)) % p

    # 计算第二个模数
    mod2 = fastMod(g, H_m, p)

    if mod1==mod2 :
        return 1
    else:
        return 0


if __name__ == "__main__":
    m = "200110727"
    print("=== 签名消息 ===")
    print("m = " + m)
    print("\n")

    PK, SK = getPKSK()
    print("=== 公钥PK和私钥SK生成结果 ===")
    print("公钥PK [p, g, y] :")
    print("p: " + str(PK[0]))
    print("g: " + str(PK[1]))
    print("y: " + str(PK[2]))
    print()
    print("私钥SK x : ")
    print("x: " + str(SK[2]))
    print("\n")

    print("=== 第1轮签名与验证(k1) ===")
    r, s = sign(m, SK)
    print("签名信息结果[r,s]为:")
    print("r: " + str(r))
    print("s: " + str(s))
    print()
    result = verify(m, r, s, PK)
    if result==0 :
        print("签名验证结果为:0，验证失败")
    else:
        print("签名验证结果为:1，验证成功")
    
    print("\n")

    print("=== 第2轮签名与验证(k2) ===")
    r, s = sign(m, SK)
    print("签名信息结果[r,s]为:")
    print("r: " + str(r))
    print("s: " + str(s))
    print()
    result = verify(m, r, s, PK)
    if result==0 :
        print("签名验证结果为:0，验证失败")
    else:
        print("签名验证结果为:1，验证成功")

    print("\n")
    print("=== 第3轮签名与验证(篡改消息m) ===")
    print("场景1假设: 双方约定的消息m被篡改【数据完整性检验】")
    print()
    print()
    r, s = sign(m, SK)
    m2 = "200110727hahaha"
    print("修改后的消息为:")
    print("m': " + m2)
    print("签名信息结果[r,s]为:")
    print("r: " + str(r))
    print("s: " + str(s))
    print()
    result = verify(m2, r, s, PK)
    if result==0 :
        print("签名验证结果为:0，验证失败")
    else:
        print("签名验证结果为:1，验证成功")
    print("\n")
    
    print("=== 第4轮签名与验证(篡改签名[r,s]) ===")
    print("场景2假设: 中间人攻击，获得[r, s]，修改了s【数据完整性检验】")
    print()
    r, s = sign(m, SK)
    print("签名信息结果[r, s]为:")
    print("r: " + str(r))
    print("s: " + str(s))
    print()
    s = s + 1
    print("篡改后的s为")
    print("s: " + str(s))
    print()
    result = verify(m, r, s, PK)
    if result==0 :
        print("签名验证结果为:0，验证失败")
    else:
        print("签名验证结果为:1，验证成功")
    print("\n")

    print("=== 第5轮签名与验证(不知道密钥SK，伪造签名) ===")
    print("场景3假设: 攻击者用自己的密钥SK来伪造签名【身份认证】")
    print()
    SK = [SK[0], SK[1], 1234567]
    print("攻击者的密钥SK为")
    print("x' :" + str(SK[2]))
    print()
    r, s = sign(m, SK)
    print("签名信息结果[r,s]为:")
    print("r: " + str(r))
    print("s: " + str(s))
    print()
    result = verify(m, r, s, PK)
    if result==0 :
        print("签名验证结果为:0，验证失败")
    else:
        print("签名验证结果为:1，验证成功")

    print("\n")